import argparse
import os
import glob
import shutil
import tqdm

parser = argparse.ArgumentParser(description='Preprocess CWE Data')


'''
    生成cpg.dot文件
'''

base_dir = '/User/zhang/data/code2graph/'

parser.add_argument('--input_dir', type=str, default='php/source-code',
                    help='The address to store the original source code directory.')
parser.add_argument('--output_dir', type=str, default='php/php-output',
                    help='The address to output the preprocessed graph.')
parser.add_argument('--joern_dir', type=str, default='/home/joern/joern-cli',
                    help='The address to joern tool')
parser.add_argument('--cuda', type=int, default=-1,
                    help='Avaiable GPU ID')


def get_all_source_file_path(input_dir):
    source_file_paths = []
    all_files = os.walk(input_dir)
    for root, dirs, files in all_files:
        for file in files:
            if file.endswith(".php") and file.startswith("CWE"):
                # print(root, dirs, file)
                source_file_path = os.path.join(root, file)
                source_file_paths.append(os.path.abspath(source_file_path))

    return source_file_paths


def get_all_source_file_dir(input_dir):
    source_file_dirs = []
    all_files = os.walk(input_dir)
    for root, dirs, files in all_files:
        for file in files:
            if file.endswith(".php") and file.startswith("CWE"):
                source_file_path = os.path.join(root, file)
                source_file_dirs.append(os.path.dirname(os.path.abspath(source_file_path)))

    return list(set(source_file_dirs))

if __name__ =="__main__":
    args = parser.parse_args()
    if not os.path.exists(args.input_dir):
        print("input data dir not exist !!!")
        exit()
    if not os.path.exists(args.joern_dir):
        print("joern tool not exist !!!")
        exit()
    if not os.path.exists(args.output_dir):
        print("output dir not exist, create output_dir")
        os.makedirs(args.output_dir)
